library(tidyverse)
library(maptools)
library(igraph)
library(parallel)
library(redist)
library(spdep)
library(gridExtra)

ipw <- function(x, beta, pop){
    indpop <- which(x$distance_parity <= pop)
    indbeta <- which(x$beta_sequence == beta)
    inds <- intersect(indpop, indbeta)
    psi <- x$constraint_pop[inds]
    w <- 1 / exp(beta * psi)
    inds <- sample(inds, length(inds), replace = TRUE, prob = w)
    x <- x$partitions[,inds]
    return(x)
}
ben_theme <- function(){
    theme_classic() +
        theme(panel.background = element_blank(),
              panel.grid.major = element_blank(),
              panel.grid.minor = element_blank(),
              axis.line = element_line(colour = "black"),
              panel.border = element_rect(colour = "black", fill = NA, size = 1),
              strip.background = element_blank(),
              legend.position = "bottom", legend.title = element_blank(),
              plot.title = element_text(hjust = 0.5),
              plot.subtitle = element_text(hjust = .5))
}

fl_map <- readShapePoly("../../data/fl/FL.shp")

## Load nodes
mn <- 89 
nodes <- strsplit(readLines("/n/imai_lab/bfifield/enumeration/largemap_enum_all/map_sub_nodes_89.dat")," ")[[1]]

## First, subset down to right map
fl_sub <- fl_map[which(rownames(fl_map@data) %in% nodes),]

## Then, reorder nodes to be correct order to match - otherwise, won't
## line up properly with original map
fl_sub <- fl_sub[match(nodes, rownames(fl_sub@data)),]

## Convert factor columns to characters, because this particular data is messy
indx <- sapply(fl_sub@data, is.factor)
fl_sub@data[indx] <- lapply(
    fl_sub@data[indx], function(x) as.numeric(as.character(x))
)

## Extract data
adjlist <- poly2nb(fl_sub, queen = FALSE)
adjlist <- lapply(adjlist, function(x){x-1})
nsims <- 100000

pop <- fl_sub@data$pop
grp <- fl_sub@data$mccain

## set seeds. 10 generated from 1-10000 uniform random.
i <- as.numeric(Sys.getenv("SLURM_ARRAY_TASK_ID"))
seeds <- c(4643, 5814, 7442, 5582, 8735, 8378, 4749, 3133, 343, 7170)
set.seed(seeds[i])
RNGkind("L'Ecuyer-CMRG")
## Begin runs

# 100%
system.time(
out_unc <-  mclapply(1:nsims, function(x){
    return(redist.rsg(
        adjlist, pop, ndists = 2, thresh = 100, verbose = FALSE,
        maxiter = 500000
    )$district_membership)
}, mc.cores = detectCores()-1,
mc.set.seed = TRUE)
)
rsg_out_unc <- do.call(cbind, out_unc)
rsg_seg_unc <- redist.segcalc(rsg_out_unc, grouppop = grp, fullpop = pop)

saveRDS(rsg_out_unc, paste0('/n/imai_lab/ckenny/enum_all/rsg_out_unc_',i,'.Rda'))
saveRDS(rsg_seg_unc, paste0('/n/imai_lab/ckenny/enum_all/rsg_seg_unc_',i,'.Rda'))

# 5% 
out_05 <-  mclapply(1:nsims, function(x){
    return(redist.rsg(
        adjlist, pop, ndists = 2, thresh = .05, verbose = FALSE,
        maxiter = 500000
    )$district_membership)
}, mc.cores = detectCores()-1,
mc.set.seed = TRUE)

rsg_out_05 <- do.call(cbind, out_05)
rsg_seg_05 <- redist.segcalc(rsg_out_05, grouppop = grp, fullpop = pop)

saveRDS(rsg_out_05, paste0('/n/imai_lab/ckenny/enum_all/rsg_out_05_',i,'.Rda'))
saveRDS(rsg_seg_05, paste0('/n/imai_lab/ckenny/enum_all/rsg_seg_05_',i,'.Rda'))

# 1%
out_01 <-  mclapply(1:nsims, function(x){
    return(redist.rsg(
        adjlist, pop, ndists = 2, thresh = .01, verbose = FALSE,
        maxiter = 500000
    )$district_membership)
}, mc.cores = detectCores()-1,
mc.set.seed = TRUE)

rsg_out_01 <- do.call(cbind, out_01)
rsg_seg_01 <- redist.segcalc(rsg_out_01, grouppop = grp, fullpop = pop)

saveRDS(rsg_out_01, paste0('/n/imai_lab/ckenny/enum_all/rsg_out_01_',i,'.Rda'))
saveRDS(rsg_seg_01, paste0('/n/imai_lab/ckenny/enum_all/rsg_seg_01_',i,'.Rda'))





